# -*- codeing = utf-8 -*-
# @Time :2021/4/3 8:59
# @Author : LUOFAN-LeiYingJie
# @File : 1-1-3-Cifar10-VGG16实现.py
# @Software : PyCharm
# @Email: 609751841@qq.com

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings

warnings.filterwarnings('ignore')
import tensorflow as tf

tf.compat.v1.logging.set_verbosity(40)

import torch
from torch import nn
import numpy as np
from torch.autograd import Variable
from torchvision.datasets import CIFAR10


class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3,out_channels=64,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True))

        self.maxpool1 = nn.MaxPool2d(kernel_size=2,padding=1)   #16 16 64

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True))

        self.maxpool2 = nn.MaxPool2d(kernel_size=2,padding=1)  #8 8 128

        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True))

        self.maxpool3 = nn.MaxPool2d(kernel_size=2)  #4 4 256

        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True))

        self.maxpool4 = nn.MaxPool2d(kernel_size=2)  #2 2 512

        self.conv5 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True))

        self.maxpool5 = nn.MaxPool2d(kernel_size=2)  #1 1 512

        self.flatten = nn.Flatten()

        self.fc1 = nn.Linear(in_features=512,out_features=4096,bias=True)
        self.relu = nn.ReLU()

        self.dropout = nn.Dropout(0.5)

        self.fc2 = nn.Linear(in_features=4096,out_features=1024,bias=True)
        self.relu = nn.ReLU()

        self.dropout = nn.Dropout(0.5)

        self.fc3 = nn.Linear(in_features=1024,out_features=10,bias=True)
        self.softmax = nn.Softmax()

    def forward(self,x):

        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.conv3(x)
        x = self.maxpool3(x)
        x = self.conv4(x)
        x = self.maxpool4(x)
        x = self.conv5(x)
        x = self.maxpool5(x)

        x = self.flatten(x)

        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)

        x = self.fc2(x)
        x = self.relu(x)
        x = self.dropout(x)

        x = self.fc3(x)
        # x = self.softmax(x)

        return x


def data_cifar10(x):
    x = np.array(x, dtype='float32') / 255.0
    x = torch.reshape(torch.from_numpy(x), shape=(32, 32, 3))
    x = np.array(x)
    x = x.transpose((2, 0, 1))  # 将channels放到第一维
    return x


def train(epoch, data):
    for step in range(epoch):

        for image, labels in data:
            # img = Variable(image.cuda())  # [bs,chaneels,H,W]
            # label = Variable(labels.cuda())  # [bs,H,W]
            img = Variable(image)  # [bs,chaneels,H,W]
            label = Variable(labels)  # [bs,H,W]

            # FP
            output = VGG16Net(img)
            loss = criterion(output, label)

            # BP
            # optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            # acc
            pre_x = torch.argmax(output, dim=1)
            num_correct = (pre_x == label).sum().item()
            acc = num_correct / img.shape[0]

            avg_loss = loss.item()

            print('Step:', step + 1, '\tloss:{:.3f}'.format(avg_loss), '\tacc:{:.2%}'.format(acc))


if __name__ == '__main__':

    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # VGG16Net = VGG16().to(device)
    VGG16Net = VGG16()

    # 定义输入为 (1, 3, 32, 32) 测试网络结构
    # input_demo = Variable(torch.zeros(1, 3, 32, 32))
    # output_demo = alexnet(input_demo)
    # print(output_demo.shape)

    train_set = CIFAR10(root='vgg/data', train=True, transform=data_cifar10, download=True)
    train_data = torch.utils.data.DataLoader(train_set, batch_size=256, shuffle=True)

    test_set = CIFAR10(root='vgg/data', train=False, transform=data_cifar10, download=True)
    test_data = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=False)

    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.Adam(VGG16Net.parameters(), lr=1e-4)

    epoch = int(input('请输入你要迭代的次数:'))
    train_name = input('请选择你要训练的数据集名称(train_data/test_data):')

    if train_name == 'train_data':
        train(epoch=epoch, data=train_data)
    else:
        train(epoch=epoch, data=test_data)